import os
import random
from PIL import Image
from  numpy.fft import fft2,fftshift
import torch
def predeal(sufffix='',ratio = 0.2):
    from shutil import copy
    
    dataroot = "/data1/yalin.huang/face_anti_spoofing_dataset/dataset1_live/HUT/Only_Face/"
    dir_D0_ori = os.path.join(dataroot, 'D0'+sufffix)  # create a path '/path/to/data/trainD0'
    dir_D45_ori = os.path.join(dataroot, 'D45'+sufffix)  # create a path '/path/to/data/trainD45'
    dir_D90_ori = os.path.join(dataroot, 'D90'+sufffix)  # create a path '/path/to/data/trainD90'
    dir_D135_ori = os.path.join(dataroot, 'D135'+sufffix)  # create a path '/path/to/data/trainD135'
    
    train_dirs = []
    dir_D0_train_dst = os.path.join(dataroot, 'D0'+sufffix,"train")  # create a path '/path/to/data/trainD0'
    dir_D45_train_dst = os.path.join(dataroot, 'D45'+sufffix,"train")  # create a path '/path/to/data/trainD45'
    dir_D90_train_dst = os.path.join(dataroot, 'D90'+sufffix,"train")  # create a path '/path/to/data/trainD90'
    dir_D135_train_dst = os.path.join(dataroot, 'D135'+sufffix,"train")  # create a path '/path/to/data/trainD135'
    train_dirs.append(dir_D0_train_dst)
    train_dirs.append(dir_D45_train_dst)
    train_dirs.append(dir_D90_train_dst)
    train_dirs.append(dir_D135_train_dst)
    
    test_dirs = []
    dir_D0_test_dst = os.path.join(dataroot, 'D0'+sufffix,"test")  # create a path '/path/to/data/testD0'
    dir_D45_test_dst = os.path.join(dataroot, 'D45'+sufffix,"test")  # create a path '/path/to/data/testD45'
    dir_D90_test_dst = os.path.join(dataroot, 'D90'+sufffix,"test")  # create a path '/path/to/data/testD90'
    dir_D135_test_dst = os.path.join(dataroot, 'D135'+sufffix,"test")  # create a path '/path/to/data/testD135'
    test_dirs.append(dir_D0_test_dst)
    test_dirs.append(dir_D45_test_dst)
    test_dirs.append(dir_D90_test_dst)
    test_dirs.append(dir_D135_test_dst)
    
    if os.path.exists(dir_D0_train_dst) is False:
        os.makedirs(dir_D0_train_dst)
        os.makedirs(dir_D45_train_dst)
        os.makedirs(dir_D90_train_dst)
        os.makedirs(dir_D135_train_dst)
        
        os.makedirs(dir_D0_test_dst)
        os.makedirs(dir_D45_test_dst)
        os.makedirs(dir_D90_test_dst)
        os.makedirs(dir_D135_test_dst)
    else:
        def del_dirsfile(dirs):
            for _dir in dirs:
                files = os.listdir(_dir)
                for file in files:
                    os.remove(os.path.join(_dir,file))
        del_dirsfile(train_dirs)
        del_dirsfile(test_dirs)
        
    
    # dirs = list([dir_D0,dir_D45,dir_D90,dir_D135])
    # for _dir in dirs:
    files = os.listdir(dir_D0_ori)
    files.remove('test')
    files.remove('train')
    files = sorted(files,key=lambda x:(int(x.split('_')[0]),int(x.split('_')[1])))
    files_dict = {}
    for file in files:
        # key = file.split('_')[0]
        key = file.split('_')[0]+'_'+file.split('_')[1]
        # if key[-1] == '2':
        #     continue
        if key in files_dict.keys():
            files_dict[key].append(file)
        else:
            files_dict[key]=[]
            files_dict[key].append(file)
            
            
    testdirs = []
    traindirs = []
    # os.makedirs()
    for key in files_dict.keys():
        random.shuffle(files_dict[key])
        for i in range(int(len(files_dict[key])*ratio)):
            testdirs.append(files_dict[key].pop())
        traindirs += files_dict[key]
        pass
    
    for _dir in testdirs:
        _dir_ori = os.path.join(dir_D0_ori,_dir)
        _dir_D45_ori = os.path.join(dir_D45_ori,_dir.split('.')[0]+'.png_D45.'+_dir.split('.')[2])
        _dir_D90_ori = os.path.join(dir_D90_ori,_dir.split('.')[0]+'.png_D90.'+_dir.split('.')[2])
        _dir_D135_ori = os.path.join(dir_D135_ori,_dir.split('.')[0]+'.png_D135.'+_dir.split('.')[2])
        
        
        _dir_dst = os.path.join(dir_D0_test_dst,_dir)
        _dir_D45_dst = os.path.join(dir_D45_test_dst,_dir.split('.')[0]+'.png_D45.'+_dir.split('.')[2])
        _dir_D90_dst = os.path.join(dir_D90_test_dst,_dir.split('.')[0]+'.png_D90.'+_dir.split('.')[2])
        _dir_D135_dst = os.path.join(dir_D135_test_dst,_dir.split('.')[0]+'.png_D135.'+_dir.split('.')[2])
        
        copy(_dir_ori,_dir_dst)
        copy(_dir_D45_ori,_dir_D45_dst)
        copy(_dir_D90_ori,_dir_D90_dst)
        copy(_dir_D135_ori,_dir_D135_dst)
    
    for _dir in traindirs:
        _dir_ori = os.path.join(dir_D0_ori,_dir)
        _dir_D45_ori = os.path.join(dir_D45_ori,_dir.split('.')[0]+'.png_D45.'+_dir.split('.')[2])
        _dir_D90_ori = os.path.join(dir_D90_ori,_dir.split('.')[0]+'.png_D90.'+_dir.split('.')[2])
        _dir_D135_ori = os.path.join(dir_D135_ori,_dir.split('.')[0]+'.png_D135.'+_dir.split('.')[2])
        
        
        _dir_dst = os.path.join(dir_D0_train_dst,_dir)
        _dir_D45_dst = os.path.join(dir_D45_train_dst,_dir.split('.')[0]+'.png_D45.'+_dir.split('.')[2])
        _dir_D90_dst = os.path.join(dir_D90_train_dst,_dir.split('.')[0]+'.png_D90.'+_dir.split('.')[2])
        _dir_D135_dst = os.path.join(dir_D135_train_dst,_dir.split('.')[0]+'.png_D135.'+_dir.split('.')[2])
        
        copy(_dir_ori,_dir_dst)
        copy(_dir_D45_ori,_dir_D45_dst)
        copy(_dir_D90_ori,_dir_D90_dst)
        copy(_dir_D135_ori,_dir_D135_dst)
        pass
    print('D0:trainLen:'+str(len(os.listdir(dir_D0_train_dst)))+'  testLen:'+str(len(os.listdir(dir_D0_test_dst))))
    print('D45:trainLen:'+str(len(os.listdir(dir_D45_train_dst)))+'  testLen:'+str(len(os.listdir(dir_D45_test_dst))))
    print('D90:trainLen:'+str(len(os.listdir(dir_D90_train_dst)))+'  testLen:'+str(len(os.listdir(dir_D90_test_dst))))
    print('D135:trainLen:'+str(len(os.listdir(dir_D135_train_dst)))+'  testLen:'+str(len(os.listdir(dir_D135_test_dst))))
        # index  = 
pass

if __name__ == '__main__':
    predeal()
    predeal('_rgb',0.2)
    predeal('_IR',0.2)
    pass
    
    
from data.image_folder import make_dataset
from data.base_dataset import BaseDataset, get_transform
from utils.util import normal
class MultiDPolarizationDataset(BaseDataset):
    """
    This dataset class can load unaligned/unpaired datasets.

    It requires two directories to host training images from domain A '/path/to/data/trainA'
    and from domain B '/path/to/data/trainB' respectively.
    You can train the model with the dataset flag '--dataroot /path/to/data'.
    Similarly, you need to prepare two directories:
    '/path/to/data/testA' and '/path/to/data/testB' during test time.
    """

    def __init__(self, opt):
        """Initialize this dataset class.

        Parameters:
            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
        """
        BaseDataset.__init__(self, opt)
        # self.dir_D0 = os.path.join(opt.dataroot, opt.phase , 'D0',opt.controller)  # create a path '/path/to/data/trainD0'
        # self.dir_D45 = os.path.join(opt.dataroot, opt.phase , 'D45',opt.controller)  # create a path '/path/to/data/trainD45'
        # self.dir_D90 = os.path.join(opt.dataroot, opt.phase , 'D90',opt.controller)  # create a path '/path/to/data/trainD90'
        # self.dir_D135 = os.path.join(opt.dataroot, opt.phase , 'D135',opt.controller)  # create a path '/path/to/data/trainD135'
        
        self.isTrain = (opt.controller == 'train')

        # self.D0_paths = sorted(make_dataset(self.dir_D0, opt.max_dataset_size))   # load images from '/path/to/data/trainA'
        # self.D45_paths = sorted(make_dataset(self.dir_D45, opt.max_dataset_size))    # load images from '/path/to/data/trainB
        # self.D90_paths = sorted(make_dataset(self.dir_D90, opt.max_dataset_size))   # load images from '/path/to/data/trainA'
        # self.D135_paths = sorted(make_dataset(self.dir_D135, opt.max_dataset_size))   # load images from '/path/to/data/trainA'
        self.D0_paths = []
        self.D45_paths = []
        self.D90_paths = []
        self.D135_paths = []
        if opt.controller == 'train':
            self.dir_D0 = os.path.join(opt.dataroot, opt.phase , 'D0'+opt._suffix,"train")  # create a path '/path/to/data/trainD0'
            self.dir_D45 = os.path.join(opt.dataroot, opt.phase , 'D45'+opt._suffix,"train")  # create a path '/path/to/data/trainD45'
            self.dir_D90 = os.path.join(opt.dataroot, opt.phase , 'D90'+opt._suffix,"train")  # create a path '/path/to/data/trainD90'
            self.dir_D135 = os.path.join(opt.dataroot, opt.phase , 'D135'+opt._suffix,"train")  # create a path '/path/to/data/trainD135'
            
            self.D0_paths += sorted(make_dataset(self.dir_D0, opt.max_dataset_size))   # load images from '/path/to/data/trainA'
            self.D45_paths += sorted(make_dataset(self.dir_D45, opt.max_dataset_size))    # load images from '/path/to/data/trainB
            self.D90_paths += sorted(make_dataset(self.dir_D90, opt.max_dataset_size))   # load images from '/path/to/data/trainA'
            self.D135_paths += sorted(make_dataset(self.dir_D135, opt.max_dataset_size))   # load images from '/path/to/data/trainA'
        elif opt.controller == 'test':
            self.dir_D0 = os.path.join(opt.dataroot, opt.phase , 'D0'+opt._suffix,"test")  # create a path '/path/to/data/trainD0'
            self.dir_D45 = os.path.join(opt.dataroot, opt.phase , 'D45'+opt._suffix,"test")  # create a path '/path/to/data/trainD45'
            self.dir_D90 = os.path.join(opt.dataroot, opt.phase , 'D90'+opt._suffix,"test")  # create a path '/path/to/data/trainD90'
            self.dir_D135 = os.path.join(opt.dataroot, opt.phase , 'D135'+opt._suffix,"test")  # create a path '/path/to/data/trainD135'
            
            self.D0_paths += sorted(make_dataset(self.dir_D0, opt.max_dataset_size))   # load images from '/path/to/data/trainA'
            self.D45_paths += sorted(make_dataset(self.dir_D45, opt.max_dataset_size))    # load images from '/path/to/data/trainB
            self.D90_paths += sorted(make_dataset(self.dir_D90, opt.max_dataset_size))   # load images from '/path/to/data/trainA'
            self.D135_paths += sorted(make_dataset(self.dir_D135, opt.max_dataset_size))   # load images from '/path/to/data/trainA'
            
            self.dir_D0 = os.path.join(opt.dataroot, opt.phase , 'D0',"train")  # create a path '/path/to/data/trainD0'
            self.dir_D45 = os.path.join(opt.dataroot, opt.phase , 'D45',"train")  # create a path '/path/to/data/trainD45'
            self.dir_D90 = os.path.join(opt.dataroot, opt.phase , 'D90',"train")  # create a path '/path/to/data/trainD90'
            self.dir_D135 = os.path.join(opt.dataroot, opt.phase , 'D135',"train")  # create a path '/path/to/data/trainD135
            
            self.D0_paths += sorted(make_dataset(self.dir_D0, opt.max_dataset_size))   # load images from '/path/to/data/trainA'
            self.D45_paths += sorted(make_dataset(self.dir_D45, opt.max_dataset_size))    # load images from '/path/to/data/trainB
            self.D90_paths += sorted(make_dataset(self.dir_D90, opt.max_dataset_size))   # load images from '/path/to/data/trainA'
            self.D135_paths += sorted(make_dataset(self.dir_D135, opt.max_dataset_size))   # load images from '/path/to/data/trainA'
        
        
        self.D0_size = len(self.D0_paths)  # get the size of dataset A
        self.D45_size = len(self.D45_paths)  # get the size of dataset B
        self.D90_size = len(self.D90_paths)  # get the size of dataset B
        self.D135_size = len(self.D135_paths)  # get the size of dataset B
        
        btoA = self.opt.direction == 'BtoA'
        input_nc = self.opt.output_nc if btoA else self.opt.input_nc       # get the number of channels of input image
        output_nc = self.opt.input_nc if btoA else self.opt.output_nc      # get the number of channels of output image
        
        self.transform_D0 = get_transform(self.opt, grayscale=(input_nc == 1),convert = self.opt.convert)
        self.transform_D45 = get_transform(self.opt, grayscale=(output_nc == 1),convert = self.opt.convert)
        self.transform_D90 = get_transform(self.opt, grayscale=(output_nc == 1),convert = self.opt.convert)
        self.transform_D135 = get_transform(self.opt, grayscale=(output_nc == 1),convert = self.opt.convert)
        self.usefft = False
        self.factor_0 = 1.0269
        self.factor_45 = 1
        self.factor_90 = 1.0385
        self.factor_135 = 1.005
        if opt.frequency:
            self.usefft=True

    def __getitem__(self, index):
        """Return a data point and its metadata information.

        Parameters:
            index (int)      -- a random integer for data indexing

        Returns a dictionary that contains A, B, A_paths and B_paths
            A (tensor)       -- an image in the input domain
            B (tensor)       -- its corresponding image in the target domain
            A_paths (str)    -- image paths
            B_paths (str)    -- image paths
        """
        D0_path = self.D0_paths[index % self.D0_size]  # make sure index is within then range
        D45_path = self.D45_paths[index]
        D90_path = self.D90_paths[index]
        D135_path = self.D135_paths[index]
        
        S0_path = D0_path.replace('/D0','/S0').replace('/train','').replace('png_D0','png_s0')
        S1_path = D0_path.replace('/D45','/S1').replace('/train','').replace('png_D45','png_s1')
        S2_path = D0_path.replace('/D90','/S2').replace('/train','').replace('png_D90','png_s2')
        # DOLp_path = D0_path.replace('/D135','/DOLP').replace('/train','').replace('png_D135','png_dolp')
        
        # if self.opt.serial_batches:   # make sure index is within then range
        #     index_D45 = index % self.D45_size
        # else:   # randomize the index for domain D45 to avoid fixed pairs.
        #     index_D45 = random.randint(0, self.D45_size - 1)
        
        # if self.opt.serial_batches:   # make sure index is within then range
        #     index_D90 = index % self.D90_size
        # else:   # randomize the index for domain D45 to avoid fixed pairs.
        #     index_D90 = random.randint(0, self.D90_size - 1)
        
        # if self.opt.serial_batches:   # make sure index is within then range
        #     index_D135 = index % self.D135_size
        # else:   # randomize the index for domain D45 to avoid fixed pairs.
        #     index_D135 = random.randint(0, self.D135_size - 1)
            
        # D45_path = self.D45_paths[index_D45]
        # D90_path = self.D90_paths[index_D90]
        # D135_path = self.D90_paths[index_D135]
        
        D0_img = Image.open(D0_path).convert('RGB')
        D45_img = Image.open(D45_path).convert('RGB')
        D90_img = Image.open(D90_path).convert('RGB')
        D135_img = Image.open(D135_path).convert('RGB')
        # apply image transformation
        D0 = self.transform_D0(D0_img)
        D45 = self.transform_D45(D45_img)
        D90 = self.transform_D90(D90_img)
        # D135 = torch.sqrt(D45**2 + D90**2)/D0
        D135 = self.transform_D135(D135_img)
        
        D0 = self.factor_0*D0
        D45 = self.factor_45*D45
        D90 = self.factor_90*D90
        D135 = self.factor_135*D135
        
        S0 =  D0 + D90
        D45_135 =  D45 + D135
        ind = torch.where(S0<D45_135)
        S0[ind] = D45_135[ind]
        
        S1 = (D0 - D90)/(S0+1e-8)
        S2= (D45 - D135)/(S0+1e-8)
        
        DOLP = torch.sqrt(S1**2 + S2**2)
        
        
        if self.usefft:
            S0_F = fftshift(fft2(S0.detach().numpy()))
            S1_F = fftshift(fft2(S1.detach().numpy()))
            S2_F = fftshift(fft2(S2.detach().numpy()))
            DOLP_F = fftshift(fft2(DOLP.cpu().detach().numpy()))

            return {'S0': S0, 'S1': S1, 'S2':S2, 'SD':DOLP,
                    'S0_paths': S0_path, 'S1_paths': S1_path, 'S2_paths': S2_path,
                    "S0_F":S0_F, "S1_F":S1_F, "S2_F":S2_F, "SD_F":DOLP_F,
                    "isTrain":self.isTrain,}
        else:
            return {'S0': S0, 'S1': S1, 'S2':S2, 'SD':DOLP,
                    'S0_paths': S0_path, 'S1_paths': S1_path, 'S2_paths': S2_path,
                    "isTrain":self.isTrain,}

    def __len__(self):
        """Return the total number of images in the dataset.

        As we have two datasets with potentially different number of images,
        we take a maximum of
        """
        return max(max(self.D0_size, self.D45_size), self.D90_size)


